using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.IO;
using System.Threading;

using Atomic.Libraries;
using Atomic.Structures;
using Atomic.Vasp;

namespace Atomic.Clusters.Fitting
{
	[Serializable]
	public class AutomaticClusterExpansionFitter
	{
		public AutomaticClusterExpansionFitter()
			: this(null)
		{
		}

		public AutomaticClusterExpansionFitter(Structure lattice)
		{
			Lattice = lattice;
			Energies = null;
			MaxClusterSize = 4;
			ClustersPerSize = new Dictionary<int, int>();
			PrintConsole = true;
			GenerateStructures = false;
		}

		public AutomaticClusterExpansionFitter(Structure lattice, IEnumerable<StructureMapping> energies)
			: this(lattice)
		{
			Energies = new ReadOnlyList<StructureMapping>(energies);
		}

		public AutomaticClusterExpansion Fit()
		{
			if (WorkingDirectory != null)
			{
				if (!WorkingDirectory.Exists)
				{
					WorkingDirectory.Create();
				}
				else if (WorkingDirectory.GetDirectories().Length > 0 || WorkingDirectory.GetFiles().Length > 0)
				{
					throw new ArgumentException("Cluster expansion directory not empty.");
				}

				return Fit(WorkingDirectory);
			}
			else
			{
				using (TempDirectory temp = new TempDirectory())
				{
					return Fit(temp.Directory);
				}
			}
		}

		private AutomaticClusterExpansion Fit(DirectoryInfo directory)
		{
			if (MaxClusterSize > 6)
			{
				throw new ArgumentException("Unsupported cluster size.");
			}
		
			Lattice.Write(directory.AddFile("lat.in"));

			if (Energies != null)
			{
				for (int i = 0; i < Energies.Count; i++)
				{
					// Use "_" to distinguish from directories created by mmaps.
					DirectoryInfo d = directory.AddDirectory("_" + i.ToString());
					d.Create();

					Structure structure = Energies[i].Structure;
					double energy = Energies[i].Value;
					structure.Write(d.AddFile("str.out"));

					// "When the calculations terminate successfully, the energy per unit cell of the structure should be copied
					// to the file 'n/energy'. (NOTE: use energy per unit cell of the structure NOT per unit cell of the lattice)."
					File.WriteAllText(d.AddFile("energy").FullName, energy.ToString("r", CultureInfo.InvariantCulture));
				}
			}

			if (ClustersPerSize.Count != 0)
			{
				foreach (int i in ClustersPerSize.Keys)
				{
					if (i < 2 || i > MaxClusterSize)
					{
						throw new ArgumentException("Invalid number of clusters requested.");
					}
				}

				using (StreamWriter sw = new StreamWriter(directory.AddFile("nbclusters.in").FullName))
				{
					for (int i = 2; i <= MaxClusterSize; i++)
					{
						int n;
						if (!ClustersPerSize.TryGetValue(i, out n))
						{
							n = 0;
						}
						sw.WriteLine(n);
					}
					sw.Close();
				}
			}

			Process p = new Process();
			try
			{
				p.StartInfo.FileName = "bash";
				p.StartInfo.WorkingDirectory = directory.FullName;
				p.StartInfo.UseShellExecute = false;
				p.StartInfo.RedirectStandardInput = true;
				if (!PrintConsole)
				{
					p.StartInfo.RedirectStandardOutput = true;
				}
				p.Start();

				// Prefer output to stdout instead of stderr.
				p.StandardInput.WriteLine(string.Format("mmaps -d -t=1 -m={0:d} -sig=17 2>&1", MaxClusterSize));
				p.StandardInput.Close();

				/*// Make sure that the ECI is fully updated. Terminate when this file is removed again. Maybe this makes mmaps perform a heavy computation twice?
				File.Create(directory.AppendFile("refresh").FullName);
				do
				{
					// Be gentle to the system; no busy waiting.
					Thread.Sleep(100);
				} while (File.Exists(directory.AppendFile("refresh").FullName));*/
				
				if (GenerateStructures)
				{
					Generate(directory);
				}

				// Request mmaps to terminate.
				directory.AddFile("stop").Create();
				p.WaitForExit();
				if (p.ExitCode != 0)
				{
					throw new Exception("Failed running mmaps.");
				}
			}
			finally
			{
				if (!p.HasExited)
				{
					p.Kill();
				}
			}

			return new AutomaticClusterExpansion(directory, Energies);
		}

		private void Generate(DirectoryInfo directory)
		{
			List<StructureMapping> currentEnergies = new List<StructureMapping>(Energies);

			// Keep a list of energies of generated structures (excluding predefined structures). The index is the integer name of the directory.
			Dictionary<int, double> energies = new Dictionary<int, double>();

			while (true)
			{
				// Ask maps to generate next structure and search for the created directory.
				directory.AddFile("ready").Create();

				while (true)
				{
					// Be gentle to the system; no busy waiting.
					Thread.Sleep(100);

					if (directory.AddFile("ready").Exists)
					{
						// Not ready yet (according to maps communication protocol).
						continue;
					}

					bool stop;
					if (ProcessNextStructure(directory, energies, currentEnergies, out stop))
					{
						// Finished processing next structure. Break and ask maps to start again.
						break;
					}

					if (stop)
					{
						return;
					}
				}

				if (directory.AddFile("stop").Exists)
				{
					// Termination was requested externally by creating this file.
					break;
				}
			}
		}

		private bool ProcessNextStructure(DirectoryInfo directory, Dictionary<int, double> energies, List<StructureMapping> currentEnergies, out bool stop)
		{
			stop = false;

			int n = 0;
			foreach (DirectoryInfo d in directory.GetDirectories())
			{
				// Look for a new integer named directory containing a wait file.
				int i;
				if (int.TryParse(d.Name, out i) && !energies.ContainsKey(i) && d.AddFile("wait").Exists)
				{
					// Compute energy for this structure using the virtual method.
					Structure structure = Structure.Read(d.AddFile("str.out"));

					double energy = NextStructure(structure, d, currentEnergies, out stop);
					currentEnergies.Add(new StructureMapping(structure, energy));

					if (stop)
					{
						return false;
					}

					if (double.IsNaN(energy) || double.IsInfinity(energy))
					{
						d.AddFile("error").Create();
					}
					else
					{
						// "When the calculations terminate successfully, the energy per unit cell of the structure should be copied
						// to the file 'n/energy'. (NOTE: use energy per unit cell of the structure NOT per unit cell of the lattice)."
						File.WriteAllText(d.AddFile("energy").FullName, energy.ToString(CultureInfo.InvariantCulture));
					}

					// Skip this next time (and remove wait according to maps communication protocol).
					energies.Add(i, energy);
					d.AddFile("wait").Delete();

					n++;
				}
			}

			// Successful if at least one structure has been processed.
			return n > 0;
		}

		private double NextStructure(Structure structure, DirectoryInfo directory, List<StructureMapping> currentEnergies, out bool stop)
		{
			double energy = double.NaN;

			AutomaticClusterExpansionNextStructureEventArgs e = new AutomaticClusterExpansionNextStructureEventArgs(structure, currentEnergies);
			if (PrepareNextStructure != null)
			{
				PrepareNextStructure(this, e);
			}

			// Skip computation if already cancelled by the event.
			if (!e.Cancel && !e.Error)
			{
				try
				{
					energy = NextStructure(e);
				}
				catch (VaspException)
				{
					// Errors from VASP are directly supported.
					e.Error = true;
				}
			}

			if (e.Error || e.Cancel)
			{
				energy = double.NaN;
			}

			stop = e.Cancel;

			return energy;
		}

		protected virtual double NextStructure(AutomaticClusterExpansionNextStructureEventArgs e)
		{
			throw new NotImplementedException("Override this method to implement.");
		}

		public event EventHandler<AutomaticClusterExpansionNextStructureEventArgs> PrepareNextStructure;

		public Structure Lattice
		{
			get;
			set;
		}

		/// <summary>
		/// The structures and corresponding energies (eV) to fit. Energies are per unit cell of the structure.
		/// </summary>
		public ReadOnlyList<StructureMapping> Energies
		{
			get;
			private set;
		}

		/// <summary>
		/// Specifies maximum clusters size in the automatic search for clusters. The default is to include up to 4-point clusters.
		/// </summary>
		public int MaxClusterSize
		{
			get;
			set;
		}

		/// <summary>
		/// Allows to manually select which clusters to include in the fit, i.e. number of pairs to include, number of triplets to include, etc.
		/// </summary>
		public Dictionary<int, int> ClustersPerSize
		{
			get;
			private set;
		}

		public DirectoryInfo WorkingDirectory
		{
			get;
			set;
		}

		public bool PrintConsole
		{
			get;
			set;
		}
		
		public bool GenerateStructures
		{
			get;
			set;
		}

		/*public List<CompositionConstraint> Compositions
		{
			get;
			private set;
		}*/
	}
}
